package com.example.caoben_back.contoller;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.content.Media;
import org.springframework.core.io.FileSystemResource;
import org.springframework.util.MimeType;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import reactor.core.publisher.Flux;
import java.io.File;
import java.net.URI;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

@RequiredArgsConstructor
@Slf4j
public class AgentWebSocketHandler extends TextWebSocketHandler {

    private final ChatClient serviceChatClient;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 1. 从WebSocket连接的URL中解析参数（prompt、chatId和fileName）
        URI uri = session.getUri();
        Map<String, String> params = null;
        if (uri != null) {
            params = parseQueryParams(uri.getQuery());
        }
        String prompt = params.get("prompt");
        String chatId = params.get("chatId");
        String fileName = params.get("fileName");

        // 2. 异步处理流式响应
        CompletableFuture.runAsync(() -> {
            try {
                // 3. 调用流式处理逻辑
                Flux<String> contentFlux;
                if (fileName != null && !fileName.isEmpty()) {
                    // 构建文件路径
                    String filePath = Paths.get("src/main/resources/agentImage", fileName).toString();
                    File imageFile = new File(filePath);
                    // 创建媒体对象
                    FileSystemResource resource = new FileSystemResource(imageFile);
                    Media media = new Media(MimeType.valueOf("image/" + getFileExtension(fileName)), resource);

                    contentFlux = serviceChatClient.prompt()
                            .options(ChatOptions.builder().model("qwen3-vl:30b").build())
                            .user(p -> p.text(prompt).media(media))
                            .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                            .stream()
                            .content();
                } else {
                    contentFlux = serviceChatClient.prompt()
                            .user(prompt)
                            .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                            .stream()
                            .content()
                            .onErrorResume(ex -> Flux.just("AI 服务异常：" + ex.getMessage()));
                }

                // 4. 将流式数据通过WebSocket发送给前端
                contentFlux.doOnComplete(() -> {
                            try {
                                session.close();
                            } catch (Exception e) {
                                log.error("Error closing session", e);
                            }
                        })
                        .subscribe(
                                content -> {
                                    try {
                                        if (session.isOpen()) {
                                            session.sendMessage(new TextMessage(content));
                                        }
                                    } catch (Exception e) {
                                        log.error("Error sending message", e);
                                    }
                                },
                                error -> log.error("Error in stream", error)
                        );
            } catch (Exception e) {
                log.error("Error processing WebSocket request", e);
                try {
                    session.sendMessage(new TextMessage("处理请求时发生错误: " + e.getMessage()));
                    session.close();
                } catch (Exception ex) {
                    log.error("Error closing session after error", ex);
                }
            }
        });
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("WebSocket connection closed with status: " + status);
    }

    // 解析URL查询参数（如?prompt=xxx&chatId=xxx&fileName=xxx）
    private Map<String, String> parseQueryParams(String query) {
        Map<String, String> params = new HashMap<>();
        if (query == null || query.isEmpty()) {
            return params;
        }
        String[] pairs = query.split("&");
        for (String pair : pairs) {
            String[] keyValue = pair.split("=");
            if (keyValue.length == 2) {
                try {
                    String key = URLDecoder.decode(keyValue[0], StandardCharsets.UTF_8);
                    String value = URLDecoder.decode(keyValue[1], StandardCharsets.UTF_8);
                    params.put(key, value);
                } catch (Exception e) {
                   log.info("Error: " + e);
                }
            }
        }
        return params;
    }

    // 获取文件扩展名
    private String getFileExtension(String fileName) {
        if (fileName == null || fileName.isEmpty()) {
            return "png"; // 默认类型
        }
        int lastDotIndex = fileName.lastIndexOf('.');
        if (lastDotIndex > 0 && lastDotIndex < fileName.length() - 1) {
            return fileName.substring(lastDotIndex + 1).toLowerCase();
        }
        return "png"; // 默认类型
    }
}